import os
import re
import sys
import zlib
import struct
import zipfile
import hashlib
import binascii
from Crypto.Cipher import AES



MAGIC_NUMBER = "OPPOENCRYPT!"
BOOT_MAGIC_NUMBER = "ANDROID!"
GZIP_MAGIC_NUMBER = "\x1F\x8B"


aes_key = {
    "r9s" : "\xD6\xDC\xCF\x0A\xD5\xAC\xD4\xE0\x29\x2E\x52\x2D\xB7\xC1\x38\x1E",
    "r9splus" : "\xD6\xDC\xCF\x0A\xD5\xAC\xD4\xE0\x29\x2E\x52\x2D\xB7\xC1\x38\x1E",
    "r11" : "\xD7\xDB\xCE\x1A\xD4\xAF\xDC\xE1\x39\x3E\x51\x21\xCB\xDC\x43\x21",
    "a77" : "\x12\x34\x1E\xAA\xC4\xC1\x23\xCE\x19\x35\x56\xA1\xBB\xCC\x23\x2D",
    #"a57" : "\x29\x23\x30\xF5\x2A\x53\x2B\x1F\xD6\xD1\xAD\xD2\x48\x3E\xC7\xE1",
    "a57" : "\xE1\xC7\x3E\x48\xD2\xAD\xD1\xD6\x1F\x2B\x53\x2A\xF5\x30\x23\x29",
    "a59st" : "",
}



def ungzip(gzip_path, ext_path):
    FTEXT, FHCRC, FEXTRA, FNAME, FCOMMENT = 1, 2, 4, 8, 16
    with open(gzip_path, "rb") as f:
        magic = f.read(2)
        if magic != '\037\213':
            return False

        method = ord(f.read(1))
        if method != 8:
            return False

        flag = ord(f.read(1))
        mtime = struct.unpack("<I", f.read(4))[0]
        f.read(2)

        if flag & FEXTRA:
            # Read & discard the extra field, if present
            xlen = ord(f.read(1))
            xlen = xlen + 256*ord(f.read(1))
            f.read(xlen)
        if flag & FNAME:
            # Read and discard a null-terminated string containing the filename
            while True:
                s = f.read(1)
                if not s or s=='\000':
                    break
        if flag & FCOMMENT:
            # Read and discard a null-terminated string containing a comment
            while True:
                s = f.read(1)
                if not s or s=='\000':
                    break
        if flag & FHCRC:
            f.read(2)     # Read & discard the 16-bit header CRC

        decompress = zlib.decompressobj(-zlib.MAX_WBITS)
        dec_data = decompress.decompress(f.read())
        if len(dec_data) == 0:
            return False

        with open(ext_path, "wb") as wf:
            wf.write(dec_data)
            return True


def sha1(path):
    with open(path, "rb") as f:
        s = hashlib.sha1()
        while True:
            data = f.read(4096)
            if len(data) == 0:
                break 
            s.update(data)
        return s.hexdigest()


def is_ozip(path):
    with open(path, "rb")  as f:
        magic_number = f.read(0x10).strip("\x00")
        if magic_number == MAGIC_NUMBER:
            return True
    return False


def ozip2zip(ozip_path, zip_path):
    with open(ozip_path, "rb") as f:
        device = ""
        f.seek(0x10, os.SEEK_SET)
        zip_size = int(f.read(0x10).strip("\x00"))
        sha1_hash = f.read(0x30).strip("\x00")
        ota_info = f.read(0x1000).strip("\x00")
        m = re.search(r"ota-id=(\S*?)_", ota_info)
        if m:
            device = m.group(0)[7:-1].lower()

        key = ""
        for k in aes_key:
            if k == device:
                key = aes_key[k]
        if key == "":
            print("can not find key")
            return False

        with open(zip_path, "wb") as wf:
            cipher = AES.new(key, AES.MODE_ECB)
            while True:
                enc_data = f.read(0x10)
                if len(enc_data) != 0x10:
                    break
                dec_data = cipher.decrypt(enc_data)
                wf.write(dec_data)

                data = f.read(0x4000)
                if len(data) == 0:
                    break
                wf.write(data)

        if sha1_hash == sha1(zip_path):
            return True

        return False


def extract_boot(zip_path, boot_path):
    with zipfile.ZipFile(zip_path, "r") as z:
        with open(boot_path, "wb") as f:
            data = z.read("boot.img")
            if len(data) == 0:
                return False
            f.write(data)
            return True


def extract_zimage(boot_path, zimage_path):
    with open(boot_path, "rb") as f:
        if f.read(8) != BOOT_MAGIC_NUMBER:
            return False
        kernel_size = struct.unpack("I", f.read(4))[0]

        f.seek(0x24, os.SEEK_SET)
        page_size = struct.unpack("I", f.read(4))[0]

        f.seek(page_size, os.SEEK_SET)
        data = f.read(kernel_size)
        if len(data) != kernel_size:
            return False
        
        with open(zimage_path, "wb") as wf:
            wf.write(data)
            return True


def extract_kernel(zimage_path, code_path):
    with open(zimage_path, "rb") as f:
        data = f.read(2)
        if data[:2] == GZIP_MAGIC_NUMBER:
            return ungzip(zimage_path, code_path)
        else:
            f.seek(0, os.SEEK_SET)
            with open(code_path, "wb") as wf:
                wf.write(f.read())


def usage():
    info = "usage: extract_boot <ozip path> [extract path]"
    print(info)


def main():
    if not len(sys.argv) > 1 and len(sys.argv) <= 3:
        usage()
        return False

    ozip_path = sys.argv[1]
    if not os.path.exists(ozip_path):
        print("ozip do not exists")
        return False

    extract_path = os.path.join(os.path.dirname(ozip_path), "extract")
    if len(sys.argv) == 3:
        extract_path = sys.argv[2]

    if not os.path.exists(extract_path):
        os.mkdir(extract_path)
    
    zip_path = os.path.join(extract_path, "rom.zip")
    boot_path = os.path.join(extract_path, "boot.img")
    zimage_path = os.path.join(extract_path, "zimage")
    code_path = os.path.join(extract_path, "code")
    if not is_ozip(ozip_path):
        print("do not ozip")
        return False

    print("ozip to zip...")
    if not ozip2zip(ozip_path, zip_path):
        print("ozip to zip fail")
        return False

    print("extract boot.img...")
    if not extract_boot(zip_path, boot_path):
        print("extract boot.img fail")
        return False

    print("extract zimage...")
    if not extract_zimage(boot_path, zimage_path):
        print("extract kernel fail")
        return False

    print("extract kernel...")
    if not extract_kernel(zimage_path, code_path):
        print("extract kernel fail")
        return False

    print("success")
    return True


if __name__ == "__main__":
    #sys.argv.append(r"D:\rom\oppo\R11Plus_11_OTA_0240_all_1WNoyo0lIwYS_local.ozip")
    #sys.argv.append(r"D:\rom")
    main()
    

